// square-root.cc
//
// Program for solving square-root step of
// the General Number Field Sieve (GNFS).
// It is based on the NTL library, which is based on GnuMP

#include <iostream>
#include <fstream>
#include <cstdlib>
#include <ctime>
#include <string>
#include <NTL/ZZ.h>
#include <NTL/RR.h>
#include <NTL/ZZX.h>
#include <NTL/ZZ_p.h>
#include <NTL/ZZXFactoring.h>
#include <NTL/ZZ_pXFactoring.h>
#include <NTL/mat_GF2.h>
#include <NTL/vec_GF2.h>
#include "nfsinput.h"
#include <limits.h>
#include <ginac/ginac.h>

// Converting from a GiNaC numeric to NTL's ZZ
NTL::ZZ ZZFromNumeric(GiNaC::numeric n) {
	NTL::ZZ res;
  	std::string s;
  	GiNaC::numeric base(256);
  	GiNaC::numeric remainder;
	GiNaC::numeric tmp=n;
	bool sign;
	
  	if(tmp.is_zero()) { 
      		return NTL::to_ZZ(0);
    	}

  	// if LONG_MIN<nx<LONG_MAX do direct conversion
  	if ( tmp<LONG_MAX && tmp>LONG_MIN ) {
		return NTL::to_ZZ(tmp.to_long()); 
    	}

  	// record sign
  	sign=tmp.is_positive();

  	// make n positive
  	if (!sign) {
		tmp = -tmp;
	}
  	// generate digits 
  	long i = 0;
  	while (tmp > 0) {
  		tmp = GiNaC::iquo(tmp, base, remainder); 
      		s= s + (char) remainder.to_int();
      		i++;
    	} 

  	NTL::ZZFromBytes(res, (const unsigned char *)s.c_str() , i);
  	if (!sign) {
		res = -res;
	}
  	return res;
}

// Converting from a NTL ZZ to GiNaC's numeric
GiNaC::numeric NumericFromZZ(NTL::ZZ n) {

	NTL::ZZ tmp=n;
	int j=NumBytes(tmp);
	unsigned char c[j];
  	int sign=1;
	GiNaC::numeric base(256);
	GiNaC::numeric exponent;
	GiNaC::ex enx=0;
	
	
	// if tmp=0 
  	if(tmp==0) { 
      		return GiNaC::numeric(0); 
    	}

	// if LONG_MIN<tmp<LONG_MAX do direct conversion
	if ( tmp<LONG_MAX && tmp>LONG_MIN ) {
      		return GiNaC::numeric(NTL::to_long(tmp)); 
    	}

  	// record sign
  	// make z positive
  	if(n<0) {
		sign=-1; 
      		tmp=-tmp;
	}

	// generate digits 
	NTL::BytesFromZZ(c, tmp, j);

	// calculate res
	for(int i=0; i<j; ++i){
    		exponent =i;
		enx += ((int)c[i])*(GiNaC::pow(base,exponent));
	}

	enx = enx*sign;
	return GiNaC::ex_to<GiNaC::numeric>(enx);
}

// Converting from a NTL ZZX to GiNaC's ex
GiNaC::ex ExFromZZX(NTL::ZZX f, GiNaC::symbol x) {

	GiNaC::ex exf;
	
	int d = NTL::deg(f);
	exf = NumericFromZZ(NTL::coeff(f,0));
	
	for(int i=0;i<d;i++){
		exf=exf+NumericFromZZ(NTL::coeff(f,i+1))*GiNaC::pow(x,i+1);
	}

	return exf;
}


// Calculates the modulus of a rational number and a integer also.
GiNaC::numeric modp(GiNaC::numeric pq, GiNaC::numeric m) {
	if(pq.is_integer()) {
		return GiNaC::mod(pq,m);
	}
	GiNaC::numeric p = pq.numer();
	GiNaC::numeric q = pq.denom();
	GiNaC::numeric p0,q0,res;

	res=p;
	q0=q;
	
	while(q0!=1) {
		p0=GiNaC::iquo(m,q0);
		
		if(GiNaC::mod(m,q0)>0) {
			p0=p0+1;
		}
		q0=GiNaC::mod(q0*p0,m);
		res*=p0;
		res=GiNaC::mod(res,m);

	}
	return res;
}

// Calculates f(x,y) mod p , f(x,y) is in Z[x][y]
// fmodp: Z[x][y] -> Z/pZ[x][y]
GiNaC::ex fmodp2(GiNaC::ex f,GiNaC::ex p,GiNaC::symbol x) {
	GiNaC::ex tmp,res,tmpf;
	tmpf=GiNaC::expand(f);
	GiNaC::ex denom=GiNaC::denom(tmpf);
	tmpf=denom*tmpf;
	
	res=tmpf.coeff(x,0);
	res=modp(GiNaC::ex_to<GiNaC::numeric>(res),GiNaC::ex_to<GiNaC::numeric>(p));
	for (int i=tmpf.ldegree(x)+1; i<=tmpf.degree(x); ++i) {
		tmp=tmpf.coeff(x,i);
		tmp=modp(GiNaC::ex_to<GiNaC::numeric>(tmp),GiNaC::ex_to<GiNaC::numeric>(p));
		if(tmp!=0)
			res=42;
	}
	return res;
}

// Calculates f(x,y) mod p , f(x,y) is in Z[x][y]
// fmodp: Z[x][y] -> Z/pZ[x][y]
GiNaC::ex fmodp(GiNaC::ex f,GiNaC::ex p,GiNaC::symbol x,GiNaC::symbol y) {
	GiNaC::ex tmp,tmpy,res;

	for (int i=f.ldegree(x); i<=f.degree(x); ++i) {
		tmp=f.coeff(x,i);
		if(GiNaC::has(tmp,y)) {
			// In worst cases we end up with a 1 degree polynomial
			// as the coefficient of x^i, we have to seperate the
			// 2 parts...
			for (int j=tmp.ldegree(y); j<=tmp.degree(y); ++j) {
				tmpy=tmp.coeff(y,j);
				tmpy=modp(GiNaC::ex_to<GiNaC::numeric>(tmpy),GiNaC::ex_to<GiNaC::numeric>(p));
				res+=tmpy*GiNaC::pow(x,i)*GiNaC::pow(y,j);
			}
		}
		else {
			tmp=modp(GiNaC::ex_to<GiNaC::numeric>(tmp),GiNaC::ex_to<GiNaC::numeric>(p));
			res+=tmp*GiNaC::pow(x,i);
		}
	}
	return res;
}

GiNaC::ex fmodp(GiNaC::ex f,GiNaC::ex p,GiNaC::symbol x) {
	GiNaC::ex tmp,res,tmpf;
	tmpf=expand(f);
	
	for (int i=tmpf.ldegree(x); i<=tmpf.degree(x); ++i) {
		tmp=tmpf.coeff(x,i);
		tmp=modp(GiNaC::ex_to<GiNaC::numeric>(tmp),GiNaC::ex_to<GiNaC::numeric>(p));
		res+=tmp*GiNaC::pow(x,i);
	}
	return res;
}


NTL::vec_GF2 makebinrep(GiNaC::ex p) {
	GiNaC::ex pnum=p;
	NTL::vec_GF2 res;
	int size=GiNaC::ex_to<GiNaC::numeric>(p).int_length();
	res.SetLength(size);

	GiNaC::ex k = GiNaC::pow(2,(size-1));
	
	for(int i=size; i>0; i--) {
		if(pnum >= k) {
			res[i-1]=1;
			pnum=pnum-k;
		}
		k=k/2;
	}
	return res;
}

inline GiNaC::ex SmodY(GiNaC::ex a,GiNaC::ex f, GiNaC::ex n, GiNaC::ex p, GiNaC::ex S, GiNaC::symbol y,GiNaC::symbol x) {
	GiNaC::symbol z("z");
	GiNaC::ex res=a;
	
	res=GiNaC::expand(res.subs(y==GiNaC::pow(z,0.5)));
	res=GiNaC::expand(res.subs(GiNaC::pow(z,0.5)==y));
	res=GiNaC::expand(res.subs(z==S));
	// reduce with f
	res=GiNaC::rem(res,f,x);
	// reduce with p
	res=fmodp(res,p,x,y);
	return res;
}
		
GiNaC::ex sqrmultmodrep(GiNaC::ex a, GiNaC::ex n, GiNaC::ex f, GiNaC::ex S, GiNaC::ex p,GiNaC::symbol y,GiNaC::symbol x ) {
	GiNaC::symbol z("z");
	GiNaC::ex res,tmpexp;
	GiNaC::ex tmpi=1;
	NTL::vec_GF2 p_bin = makebinrep(n);
	
	res=a;
	for(int i=0; i < p_bin.length(); i++) {
		if(p_bin[i]==1) {
			tmpi=tmpi*res;
			tmpi=SmodY(tmpi,f,n,p,S,y,x);
		}
		res=res*res;
		res=SmodY(res,f,n,p,S,y,x);
		//std::cout << i << " " << tmpi << std::endl;
	}
	return tmpi;
}

GiNaC::ex AdjustPoly(GiNaC::ex f, GiNaC::ex p, GiNaC::symbol x) {
	GiNaC::ex res=0;
	GiNaC::ex tmpres;
	
	for (int i=f.ldegree(x); i<=f.degree(x); ++i) {
		tmpres=f.coeff(x,i);
		
		if(tmpres > p/2) { 
			tmpres-=p;
		}
		res+=tmpres*GiNaC::pow(x,i);
	}

	return res;
}


/////////////////////////////////////////////////////////// 
// Calculates the squareroot of S using Buhler's method  //
///////////////////////////////////////////////////////////
GiNaC::ex sqrtmodp(GiNaC::ex R, GiNaC::ex S, GiNaC::ex f, GiNaC::ex P,GiNaC::symbol x ,GiNaC::symbol y) {
	GiNaC::ex res,tmpexp,tmpres,RS,prev;
	GiNaC::numeric n=1;
	res=R;
	
	while(true) {
		tmpres=res;
		tmpexp=GiNaC::pow(P,2*n);
		
		res=GiNaC::pow(tmpres,2);
		res=(S*res);
		res=GiNaC::rem(res,f,x);
		res=fmodp(res,tmpexp,x);
		
		res=(tmpres*(3-res))/2;
		res=GiNaC::rem(res,f,x);
		res=fmodp(res,tmpexp,x);
			
		tmpres=S*res;
		tmpres=GiNaC::rem(tmpres,f,x);
		RS=fmodp(tmpres,tmpexp,x);
		tmpres=GiNaC::pow(AdjustPoly(RS,tmpexp,x),2);
		tmpres=GiNaC::rem(tmpres,f,x);
		
		if(tmpres==S) {
			break;
		}
		if(n>7)
			break;
		n++;
	}
	return AdjustPoly(RS,tmpexp,x);
	
}

// Calculates random prime p larger than n so
// that f mod p is irreducible over Z
GiNaC::ex RandomPrime(GiNaC::ex n, NTL::ZZX f) {
	
	int d = NTL::deg(f);
	int pbit=3*GiNaC::ex_to<GiNaC::numeric>(n).int_length();
	NTL::ZZ q=NTL::RandomPrime_ZZ(pbit,50);
	NTL::ZZ_p::init(q);

	NTL::ZZ_pX fp;
	NTL::ZZ_p ap;
	
	while(true) {
		for(int i=0;i<d;i++){
			ap=NTL::to_ZZ_p(NTL::coeff(f,i+1));
			NTL::SetCoeff(fp,i,ap);
		}
		if(NTL::DetIrredTest(fp)) {
			break;
		}
		q=NTL::RandomPrime_ZZ(pbit,50);
		NTL::ZZ_p::init(q);
	}
	
	std::cout << "Chosen new " << pbit << " bit prime... " << std::endl;
	return NumericFromZZ(q);
}

// Calculates random polynomial in x of degree d with 
// coefficients mod p i.e. r \in Z_p[x]
GiNaC::ex RandomPoly(GiNaC::ex d,GiNaC::ex p, GiNaC::symbol x) {
	
	GiNaC::ex r=0;
	GiNaC::ex a,i;
	NTL::ZZ ntl_p=ZZFromNumeric(GiNaC::ex_to<GiNaC::numeric>(p));
	for(i=0; i<=d;i++) {
		a = NumericFromZZ(NTL::RandomBnd(ntl_p));
		r+=a*pow(x,i);
	}
	return r;
}

// Calculates the squareroot of the rational number
// f'(m)^2*s(m) and evaluates it mod n
GiNaC::ex EvalRationalSqrt(GiNaC::ex f,GiNaC::ex s,GiNaC::ex m,GiNaC::ex n,GiNaC::symbol x) {

	GiNaC::ex res,root1,root2;
	GiNaC::ex df=f.diff(x);
	res = GiNaC::sqrt(s.subs(x==m))*df.subs(x==m);
	return modp(GiNaC::ex_to<GiNaC::numeric>(res),GiNaC::ex_to<GiNaC::numeric>(n));
}

// Calculates the squareroot of the algebraic number
// f'(x)^2*s and evaluates it in m mod n
GiNaC::ex EvalAlgebraicSqrt(GiNaC::ex f,GiNaC::ex s,GiNaC::ex m,GiNaC::ex n,GiNaC::symbol x, GiNaC::ex pex,NTL::ZZX fzz) {
	GiNaC::symbol y("y");
	GiNaC::ex tmp=0;
	GiNaC::ex df=f.diff(x);
	GiNaC::ex r,r1,r2,q,algroot;
	GiNaC::ex denominator=1;	
	
	// See if S/f(x) is in Q[x]
	// if it is - remove common denominator
	// and save it for later....
	GiNaC::ex smodp=GiNaC::rem(s,f,x);
	denominator=GiNaC::denom(smodp);
	GiNaC::ex smodp2=smodp*denominator;
	GiNaC::ex p=RandomPrime(n,fzz);
	GiNaC::ex pexp = (p-1)/2;
	int k=0;
	int o=0;
	while(true) {
		k=0;
		while(true) {
			r1= sqrmultmodrep(y-o,pexp,f,smodp2,p,y,x);
			r1 = r1.coeff(y,1);
				
			o++;
			// check for usability !!
			if(r1!=0) {
				tmp=smodp2*GiNaC::pow(r1,2);
				tmp=GiNaC::rem(tmp,f,x);
				tmp=fmodp(tmp,p,x,y);
				if (tmp == 1) {
					break;
				}
			}
			if(k>30) {
				p=RandomPrime(n,fzz);
				pexp = (p-1)/2;
				k=0;
				o=0;
			}
			k++;
		}
	
		tmp=sqrtmodp(r1,smodp2,f,p,x,y);
		tmp=tmp.subs(x==m)*(df.subs(x==m));
		tmp=modp(GiNaC::ex_to<GiNaC::numeric>(tmp),GiNaC::ex_to<GiNaC::numeric>(n));
		tmp=GiNaC::sqrt(1/denominator)*tmp;
		algroot=modp(GiNaC::ex_to<GiNaC::numeric>(tmp),GiNaC::ex_to<GiNaC::numeric>(n));
		q=modp(GiNaC::ex_to<GiNaC::numeric>(GiNaC::pow(algroot,2)),GiNaC::ex_to<GiNaC::numeric>(n));
		if(pex==q) {	
			std::cout << "\t\t\t\t\t\t :-) " << std::endl;
			return algroot;
		}
		else{
			k= -200;
		}
	}
	return algroot;
}


NTL::vec_GF2 getSolution(NTL::mat_GF2* M, NTL::vec_GF2* freecols, int num) {
	// finds a solution from the reduced matrix and the 
	// vector of free variables...
	// it takes a number representing the number
	// of solution wanted.

	int j=-1;
	int i=num;
	int h;
	NTL::vec_GF2 res;
	res.SetLength(freecols->length());
	while(i>0) {
		j++;
		while((*freecols)[j]==0) {
			j++;
		}
		i--;
	}
	res[j]=1;
	for(i=0;i<M->NumRows();i++) {
		if((*M)[i][j]==1) {
			h=i;
			while(h<j) {
				if((*M)[i][h]==1) {
					res[h]=1;
					break;
				}
				h++;
			}
		}
	}
	return res;
}
GiNaC::ex DepsfileToEx(char* depfile,GiNaC::symbol x) {
	GiNaC::ex s=1;
	NTL::ZZ a,b;
	int i=0;	
	
	std::ifstream infile(depfile);

	while(!infile.eof()) {
		infile >> a;
		infile >> b;
		s=s*(NumericFromZZ(a)+NumericFromZZ(b)*x);
		i++;
	}
	
	infile.close();
	std::cout << "number of relations in solution: " << i << std::endl;
	return s;
}

void findDeps(NTL::vec_GF2 rel, char* relfile, char* depfile) {
	// reads the dependecies from the relations-file and
	// write them to depfile
	int i=0;
	NTL::ZZ a,b;
	bool first = true;

	std::ifstream infile(relfile);
	std::ofstream outfile(depfile);
	while(i<rel.length()) {
		infile >> a;
		infile >> b;

		if(rel[i]==1) {
			if(first) {
				outfile << a << std::endl << b;
				first = false;
			}
			else {
				outfile << std::endl << a << std::endl << b ;
			}
		}
		i++;
	}	
	outfile.close();
	std::cout << "Dependencies stored in: " << depfile << std::endl;
	infile.close();
}

NTL::vec_ZZ squareRoot(NTL::ZZ n,nfsPoly poly, NTL::mat_GF2* M, NTL::vec_GF2* freecols,char* relfile, char* depfile) {

	NTL::SetSeed(NTL::to_ZZ(time(NULL)));
	int i=1;
	NTL::vec_GF2 solution;
	NTL::vec_ZZ res;
	NTL::ZZ gcd;
	bool done=false;
	GiNaC::symbol x("x");
	GiNaC::ex s,algroot,ratroot,p,q;
	GiNaC::ex nex = NumericFromZZ(n);	
	GiNaC::ex mex = NumericFromZZ(poly.m);
	GiNaC::ex fex = ExFromZZX(poly.f,x);
	GiNaC::ex pex;

	std::cout << "n: " << nex << std::endl;
	std::cout << "m: " << mex << std::endl;
	std::cout << "f: " << fex << std::endl;

	// number of free variables
	int freevals=0;
	for(int k=0;k<freecols->length();k++) {
		if((*freecols)[k]==1) {
			freevals++;
		}
	}
	std::cout << "number of free variables: " << freevals << std::endl;
	// trying to calculate squareroots...
	while (i<freevals && !done) {
		solution=getSolution(M,freecols,i);
		findDeps(solution,relfile,depfile);
		s=DepsfileToEx(depfile,x);
		ratroot = EvalRationalSqrt(fex,s,mex,nex,x);
		std::cout << "Rational root: " << ratroot << std::endl;
		p=modp(GiNaC::ex_to<GiNaC::numeric>(GiNaC::pow(ratroot,2)),GiNaC::ex_to<GiNaC::numeric>(nex));
		algroot = EvalAlgebraicSqrt(fex,s,mex,nex,x,p,poly.f);
		std::cout << "Algebraic root: " << algroot << std::endl;
		q=modp(GiNaC::ex_to<GiNaC::numeric>(GiNaC::pow(algroot,2)),GiNaC::ex_to<GiNaC::numeric>(nex));
	
		if(p==q) {
			p=GiNaC::gcd(nex,algroot-ratroot);
			q=GiNaC::gcd(nex,algroot+ratroot);

			if ((p==1 || p==nex) && (q==1 || q==nex)) {
				std::cout << "Found trivial divisors :-(" 
					  << std::endl;
			}
			if ((p!=1 && p!=nex) || (q!=1 && q!=nex)) {
				std::cout << "divisors of " << n 
					  << " is: " << p << " " << q 
					  << std::endl;
				done=true;
			}
		}
		if(!done) {
			std::cout << "Trying different set of relations...\n";
		}
		i++;
	}
	return res;	
}


int main (int count, char** args){
	if(count==1) {
		std::cout << std::endl;
		std::cout << "Usage: ./pgnfs_sqrt init-file"  << std::endl;
		std::cout << std::endl;
		return 1;
	}
	std::clock_t start,end;
	
	NFS_instance prob = readData(args[1]);
	std::string polyfile=args[1];
	polyfile+=".poly";
	nfsPoly poly = readPoly((char*) polyfile.c_str());

	std::string rationalfile=args[1];
	rationalfile+=".rfb";
	std::string algebraicfile=args[1];
	algebraicfile+=".afb";
	std::string quadraticfile=args[1];
	quadraticfile+=".qcb";
	std::string relationsfile=args[1];
	relationsfile+=".rel";
	std::string matrixfile=args[1];
	matrixfile+=".mat";
	std::string depsfile=args[1];
	depsfile+=".dep";

	NTL::RR::SetPrecision(NTL::NumBits(prob.n));
	NTL::mat_GF2 matrix;
	NTL::vec_GF2 freecols;
	NTL::vec_ZZ res;

	std::ifstream matfile((char*) matrixfile.c_str());
	matfile >> matrix;
	matfile >> freecols;
	matfile.close();
	NTL::mat_GF2* M = &matrix;
	NTL::vec_GF2* freevars = &freecols;
	
	start=std::clock();

	std::cout << "Calculating squareroot..." << std::endl;
	res=squareRoot(prob.n,poly,M, freevars,(char*) relationsfile.c_str(),(char*) depsfile.c_str());
	end=std::clock();	
	std::cout << "Squareroot calculated...in " <<(double) (end - start)/CLOCKS_PER_SEC <<  " seconds" << std::endl;
}

